#!/usr/bin/env julia

"""
# Comprehensive Tests for Advanced Solver Capabilities

Tests for the research-grade solver implementations including:
- Newton-Krylov methods
- Multigrid preconditioning  
- Adaptive time stepping
- Block preconditioning methods
"""

using Test
using NSEMSolver
using LinearAlgebra
using SparseArrays
using Random

# Set random seed for reproducible tests
Random.seed!(42)

@testset "Advanced Solvers Test Suite" begin
    
    # ============================================================================
    # Newton-Krylov Methods Tests
    # ============================================================================
    
    @testset "Newton-Krylov Methods" begin
        
        @testset "NewtonKrylovSolver Construction" begin
            # Test solver creation
            A = spdiagm(-1 => -ones(9), 0 => 2*ones(10), 1 => -ones(9))
            linear_solver = JuliaNSSolver(A, :gmres; max_iter=100, tolerance=1e-8)
            
            nk_solver = NewtonKrylovSolver(
                linear_solver=linear_solver,
                tolerance=1e-6,
                max_iterations=10,
                linesearch=:backtrack,
                jacobian_method=:fd
            )
            
            @test nk_solver.tolerance == 1e-6
            @test nk_solver.max_iterations == 10
            @test nk_solver.linesearch == :backtrack
            @test nk_solver.jacobian_method == :fd
        end
        
        @testset "Linear Problem via Newton-Krylov" begin
            # Test Newton-Krylov on a linear problem (should converge in 1-2 iterations)
            n = 20
            A = spdiagm(-1 => -ones(n-1), 0 => 3*ones(n), 1 => -ones(n-1))
            b = ones(n)
            
            function linear_residual!(residual, x)
                residual .= A * x .- b
                return nothing
            end
            
            linear_solver = JuliaNSSolver(A, :gmres; max_iter=50, tolerance=1e-10)
            nk_solver = NewtonKrylovSolver(linear_solver=linear_solver, tolerance=1e-8, 
                                         max_iterations=5, verbose=false)
            
            x0 = zeros(n)
            result = solve_nonlinear_ns!(nk_solver, linear_residual!, nothing, x0)
            
            @test result.converged == true
            @test result.iterations <= 3  # Should converge very quickly for linear problem
            @test result.final_residual < 1e-6
        end
        
        @testset "Jacobian-vector Product" begin
            # Test finite difference Jacobian-vector product
            n = 10
            A = spdiagm(0 => ones(n))
            
            function test_residual!(residual, x)
                residual .= x.^2  # Simple nonlinear function
                return nothing
            end
            
            linear_solver = JuliaNSSolver(A, :gmres)
            nk_solver = NewtonKrylovSolver(linear_solver=linear_solver, jacobian_epsilon=1e-6)
            
            x = ones(n)
            v = ones(n)
            jv = compute_jacobian_action!(nk_solver, test_residual!, x, v)
            
            # For f(x) = x^2, Jacobian is 2*x, so J*v = 2*x*v = 2*ones
            expected_jv = 2 * ones(n)
            @test norm(jv - expected_jv) < 1e-4
        end
    end
    
    # ============================================================================
    # Multigrid Preconditioning Tests
    # ============================================================================
    
    @testset "Multigrid Preconditioning" begin
        
        @testset "Multigrid Hierarchy Creation" begin
            # Test multigrid level creation
            opts = NSOptions(N=4, n_block=2, dim=2)
            multidomain = create_multidomain(opts.n_block, opts.N, 0.01, opts.dim)
            sem_operators = create_sem_operators(multidomain)
            
            levels = create_multigrid_hierarchy(sem_operators, multidomain, 3;
                                              coarsening_strategy=:p_multigrid,
                                              min_polynomial_order=2)
            
            @test length(levels) >= 2
            @test length(levels) <= 3
            
            # Check polynomial order decreases
            if length(levels) > 1
                @test levels[1].polynomial_order > levels[2].polynomial_order
            end
        end
        
        @testset "Restriction/Prolongation Operators" begin
            # Test p-multigrid operators
            fine_order = 4
            coarse_order = 2
            
            R = create_p_restriction_operator(fine_order, coarse_order)
            P = create_p_prolongation_operator(coarse_order, fine_order)
            
            @test size(R, 1) == (coarse_order + 1)^2  # 2D case
            @test size(R, 2) == (fine_order + 1)^2
            @test size(P, 1) == (fine_order + 1)^2
            @test size(P, 2) == (coarse_order + 1)^2
            
            # Test basic properties
            @test nnz(R) > 0
            @test nnz(P) > 0
        end
        
        @testset "Multigrid Preconditioner Setup" begin
            # Test multigrid preconditioner creation
            n = 100
            A = spdiagm(-1 => -ones(n-1), 0 => 2*ones(n), 1 => -ones(n-1))
            
            # Create minimal multigrid setup
            levels = [
                MultigridLevel(nothing, nothing, 4, 4, nothing, nothing, A, nothing),
                MultigridLevel(nothing, nothing, 2, 4, nothing, nothing, A[1:50, 1:50], nothing)
            ]
            
            mg_preconditioner = MultigridPreconditioner(levels=levels, verbose=false)
            
            @test length(mg_preconditioner.levels) == 2
            @test mg_preconditioner.smoother == :jacobi
            @test mg_preconditioner.cycle_type == :V
        end
    end
    
    # ============================================================================
    # Adaptive Time Stepping Tests  
    # ============================================================================
    
    @testset "Adaptive Time Stepping" begin
        
        @testset "AdaptiveTimestepper Construction" begin
            stepper = AdaptiveTimestepper(
                method=:rk45,
                tolerance=1e-4,
                safety_factor=0.9
            )
            
            @test stepper.method == :rk45
            @test stepper.tolerance == 1e-4
            @test stepper.safety_factor == 0.9
        end
        
        @testset "Butcher Tableau Retrieval" begin
            # Test different RK methods
            tableau_rk23 = get_butcher_tableau(:rk23)
            tableau_rk45 = get_butcher_tableau(:rk45) 
            tableau_dopri5 = get_butcher_tableau(:dopri5)
            
            @test tableau_rk23.order == 2
            @test tableau_rk23.error_order == 3
            @test size(tableau_rk23.A, 1) == length(tableau_rk23.c)
            
            @test tableau_rk45.order == 5
            @test tableau_rk45.error_order == 4
            
            @test tableau_dopri5.order == 5
            @test tableau_dopri5.error_order == 4
        end
        
        @testset "Error Estimation" begin
            # Test error estimation function
            n = 10
            u_high = ones(n)
            u_low = 1.01 * ones(n)  # Small difference
            v_high = zeros(n)
            v_low = 0.01 * ones(n)  # Small difference
            
            error_est = estimate_error(u_high, u_low, v_high, v_low, 1e-6, 1e-6)
            
            @test error_est > 0
            @test error_est < 1  # Should be reasonable
        end
        
        @testset "Step Size Update" begin
            stepper = AdaptiveTimestepper(pi_controller=true, safety_factor=0.9)
            
            # Test step size increase for small error
            dt_new, factor = update_timestep(0.01, 0.1, stepper)  # Error < 1, should increase
            @test dt_new >= 0.01
            @test factor >= 1.0
            
            # Test step size decrease for large error  
            dt_new, factor = update_timestep(0.01, 2.0, stepper)  # Error > 1, should decrease
            @test dt_new <= 0.01
            @test factor <= 1.0
        end
        
        @testset "Stability Monitoring" begin
            # Test stability monitoring
            opts = NSOptions(N=4, n_block=2)
            
            # Stable solution
            u_stable = 0.5 * ones(10, 10)
            v_stable = 0.3 * ones(10, 10)
            @test monitor_stability(u_stable, v_stable, 0.01, opts) == true
            
            # Unstable solution (NaN)
            u_unstable = fill(NaN, 10, 10)
            v_unstable = zeros(10, 10)
            @test monitor_stability(u_unstable, v_unstable, 0.01, opts) == false
            
            # Unstable solution (too large)
            u_large = 1e8 * ones(10, 10)
            v_large = zeros(10, 10)
            @test monitor_stability(u_large, v_large, 0.01, opts) == false
        end
    end
    
    # ============================================================================
    # Block Preconditioning Tests
    # ============================================================================
    
    @testset "Block Preconditioning" begin
        
        @testset "Schur Complement Solver Construction" begin
            schur_solver = create_schur_complement_solver(:julia, :julia;
                                                         schur_approximation=:pressure_mass)
            
            @test isa(schur_solver, SchurComplementSolver)
            @test schur_solver.schur_approximation == :pressure_mass
            @test isa(schur_solver.velocity_solver, JuliaNSSolver)
            @test isa(schur_solver.pressure_solver, JuliaNSSolver)
        end
        
        @testset "Block Triangular Solver Construction" begin
            block_solver = create_block_triangular_solver(:julia, :julia;
                                                         approximation_type=:upper)
            
            @test isa(block_solver, BlockTriangularSolver)
            @test block_solver.approximation_type == :upper
        end
        
        @testset "Pressure Mass Matrix Creation" begin
            grid_info = (n=3, n_block=2, dim=2, index_sum=4)
            polynomial_order = 3
            
            M = create_pressure_mass_matrix(grid_info, polynomial_order)
            
            @test isa(M, AbstractMatrix)
            @test size(M, 1) == size(M, 2)  # Square matrix
            @test all(diag(M) .> 0)  # Positive diagonal entries
        end
        
        @testset "Simple Schur Complement Solve" begin
            # Create small test system
            n_vel = 20
            n_pres = 10
            
            A = spdiagm(-1 => -ones(n_vel-1), 0 => 3*ones(n_vel), 1 => -ones(n_vel-1))
            B = sprandn(n_pres, n_vel, 0.2)  # Random sparse divergence operator
            
            f = randn(n_vel)
            g = randn(n_pres)
            rhs = [f; g]
            
            # Create Schur solver
            schur_solver = create_schur_complement_solver(:julia, :julia; verbose=false)
            
            # Set up solver
            system_matrices = (velocity_matrix=A, divergence_matrix=B)
            grid_info = (n=3, n_block=2, dim=2, index_sum=4)
            setup_block_preconditioner!(schur_solver, system_matrices, grid_info)
            
            initial_guess = zeros(length(rhs))
            
            # This test mainly checks that the function runs without error
            # Full accuracy testing would require a more sophisticated setup
            @test_nowarn begin
                solution = solve_schur_complement!(schur_solver, rhs, initial_guess, 
                                                 A, B, :coupled, false)
            end
        end
    end
    
    # ============================================================================
    # Integration Tests
    # ============================================================================
    
    @testset "Solver Integration Tests" begin
        
        @testset "Compute Stable Timestep" begin
            opts = NSOptions(N=4, n_block=2, nu=0.01)
            multidomain = create_multidomain(opts.n_block, opts.N, 0.01, opts.dim)
            
            u = 0.5 * ones(10, 10)
            v = 0.3 * ones(10, 10)
            
            dt_stable = compute_stable_timestep(u, v, multidomain, opts)
            
            @test dt_stable > 0
            @test dt_stable < 1.0  # Should be reasonable
            @test isfinite(dt_stable)
        end
        
        @testset "Enhanced Time Stepping" begin
            opts = NSOptions(N=3, n_block=2)
            
            u = ones(5, 5)
            v = zeros(5, 5)
            u_new = similar(u)
            v_new = similar(v)
            
            # Test enhanced RK4 implementation
            @test_nowarn runge_kutta_step!(u_new, v_new, u, v, 0.01, opts)
            
            # Check that solution changed
            @test u_new != u  # Should be different due to RK4 update
            @test v_new != v
        end
    end
    
    # ============================================================================
    # Utility Function Tests
    # ============================================================================
    
    @testset "Utility Functions" begin
        
        @testset "DOF Estimation" begin
            grid_info = (n=4, n_block=3, dim=2, index_sum=9)
            polynomial_order = 4
            
            n_dofs = estimate_pressure_dofs(grid_info, polynomial_order)
            
            @test n_dofs > 0
            @test n_dofs == grid_info.index_sum * (polynomial_order + 1)^grid_info.dim
        end
        
        @testset "Create Adaptive Timestepper" begin
            stepper = create_adaptive_timestepper(:dopri5; tolerance=1e-5, verbose=true)
            
            @test stepper.method == :dopri5
            @test stepper.tolerance == 1e-5
            @test stepper.verbose == true
        end
    end
end

# ============================================================================
# Performance and Stress Tests
# ============================================================================

@testset "Performance Tests" begin
    
    @testset "Newton-Krylov Performance" begin
        # Test that Newton-Krylov can handle medium-size problems efficiently
        n = 50
        A = spdiagm(-1 => -ones(n-1), 0 => 3*ones(n), 1 => -ones(n-1))
        b = ones(n)
        
        function residual_fn!(residual, x)
            residual .= A * x .- b .+ 0.01 .* x.^3  # Mild nonlinearity
            return nothing
        end
        
        linear_solver = JuliaNSSolver(A, :gmres; max_iter=100, tolerance=1e-8)
        nk_solver = NewtonKrylovSolver(linear_solver=linear_solver, tolerance=1e-6, 
                                     max_iterations=10, verbose=false)
        
        x0 = zeros(n)
        
        # Time the solve
        start_time = time()
        result = solve_nonlinear_ns!(nk_solver, residual_fn!, nothing, x0)
        solve_time = time() - start_time
        
        @test result.converged == true
        @test solve_time < 2.0  # Should solve reasonably quickly
        @test result.solve_time < 2.0
    end
    
    @testset "Multigrid Overhead" begin
        # Test that multigrid setup doesn't have excessive overhead
        opts = NSOptions(N=3, n_block=2, dim=2)
        multidomain = create_multidomain(opts.n_block, opts.N, 0.01, opts.dim)
        sem_operators = create_sem_operators(multidomain)
        
        start_time = time()
        levels = create_multigrid_hierarchy(sem_operators, multidomain, 3)
        creation_time = time() - start_time
        
        @test creation_time < 1.0  # Should create hierarchy quickly
        @test length(levels) >= 2
    end
end

println("\n✅ All advanced solver tests completed successfully!")
println("🚀 NSEMSolver.jl advanced capabilities are verified and ready for research use!")